# -*- coding: utf-8 -*-

"""
Created on 03/23/2022
train_test.
@author: Kang Xiatao (kangxiatao@gmail.com)
"""

import os

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from utils.common_utils import PresetLRScheduler
from regul import *
from pruner.pruning_exp import fetch_data
import torch.autograd as autograd
import copy


def conv2d_norm(model):
    l2_norm = []
    for module in model.modules():
        if type(module) is nn.Conv2d:
            l2_norm.append((module.weight ** 2).sum())
    return sum(l2_norm)**0.5


# def l2_regularization(model, weight_decay=0.0005):
#     l2_loss = 0
#     for layer in model.modules():
#         if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
#             l2_loss += (layer.weight ** 2).sum()
#     return weight_decay * l2_loss / 2.0


# def l2_regularization(model, l2_alpha=0.0005):
#     for module in model.modules():
#         if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
#             module.weight.grad.data.add_(l2_alpha * module.weight.data)


def train(net, loader, optimizer, criterion, epoch, writer, lr_mode, reg=None):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0

    desc = ('- Loss: %.3f | Acc: %.3f%% (%d/%d)' % (0, 0, correct, total))

    # prog_bar = tqdm(enumerate(loader), total=len(loader), desc=desc, leave=True)
    # for batch_idx, (inputs, targets) in prog_bar:
    for batch_idx, (inputs, targets) in enumerate(loader):
        inputs, targets = inputs.cuda(), targets.cuda()
        optimizer.zero_grad()
        outputs = net(inputs)

        loss = criterion(outputs, targets) + reg(net)
        # loss = criterion(outputs, targets)
        loss.backward()
        # reg(net)
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        # desc = ('- Loss: %.3f | Acc: %.3f%% (%d/%d)' % (train_loss/(batch_idx+1), 100.*correct/total, correct, total))
        # prog_bar.set_description(desc, refresh=True)
    print('- Loss: %.3f | Acc: %.3f%% (%d/%d)' % (train_loss/(batch_idx+1), 100.*correct/total, correct, total))

    if writer:
        writer.add_scalar('train/loss', train_loss / (batch_idx + 1), epoch)
        writer.add_scalar('train/acc', 100. * correct / total, epoch)

        # 卷积范数值
        writer.add_scalar('train/conv2d_norm', conv2d_norm(net), epoch)


def test(net, loader, criterion, epoch, writer):
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    desc = ('Loss: %.3f | Acc: %.3f%% (%d/%d)'
            % (test_loss / (0 + 1), 0, correct, total))

    # prog_bar = tqdm(enumerate(loader), total=len(loader), desc=desc, leave=True)
    with torch.no_grad():
        # for batch_idx, (inputs, targets) in prog_bar:
        for batch_idx, (inputs, targets) in enumerate(loader):
            inputs, targets = inputs.cuda(), targets.cuda()
            outputs = net(inputs)
            loss = criterion(outputs, targets)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

            # desc = ('Loss: %.3f | Acc: %.3f%% (%d/%d)'
            #         % (test_loss / (batch_idx + 1), 100. * correct / total, correct, total))
            # prog_bar.set_description(desc, refresh=True)

    acc = 100. * correct / total
    print('Loss: %.3f | Acc: %.3f%% (%d/%d)' % (test_loss/(batch_idx+1), acc, correct, total))

    if writer:
        writer.add_scalar('test/loss', test_loss / (batch_idx + 1), epoch)
        writer.add_scalar('test/acc', 100. * correct / total, epoch)
    return acc


def writer_exp(net, loader, criterion, writer, epoch, num_classes, exp_name, train_mode=False):
    if train_mode:
        net.train()
    else:
        net.eval()
    # 重新计算损失和建图，便于求一阶二阶导
    weights = []
    for layer in net.modules():
        if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
            weights.append(layer.weight)
    for w in weights:
        w.requires_grad_(True)

    inputs, targets = fetch_data(loader, 1, num_classes, 1)  # 随机取100个样本 random
    inputs = inputs.cuda()
    targets = targets.cuda()

    # 梯度耦合
    num_group = 2
    N = inputs.shape[0]
    equal_parts = N // num_group
    grad_ls = []
    gss_ls = []
    for i in range(num_group):
        _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
        _loss = criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts])
        grad_ls.append(autograd.grad(_loss, weights, create_graph=True))

    for i in range(num_group):
        for j in range(i + 1, num_group):
            _gz = 0
            _gz_hi = 0
            _gz_hj = 0
            _layer = 0
            for layer in net.modules():
                if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
                    _gz += (grad_ls[i][_layer] * grad_ls[j][_layer]).sum()  # g1 * g2
                    _layer += 1
            gss_ls.append(_gz)

    # 梯度
    _outputs = net.forward(inputs)
    _loss = criterion(_outputs, targets)
    _g = autograd.grad(_loss, weights, create_graph=True)
    gtg = 0
    _layer = 0
    for layer in net.modules():
        if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
            gtg += _g[_layer].pow(2).sum()  # g * g
            _layer += 1

    # tensorboard
    if train_mode:
        exp_name = 'train_' + exp_name
    writer.add_scalar(exp_name + '/gtg', gtg, epoch)
    writer.add_scalar(exp_name + '/gss', gss_ls[0], epoch)


def loss_coupling(net, loader, criterion, scheduler, epoch, writer, samples_num=256):

    net = copy.deepcopy(net).train()
    net.zero_grad()
    _optim = optim.SGD(net.parameters(), lr=scheduler.get_last_lr()[0], momentum=0.9, weight_decay=0.0005)

    inputs, targets = fetch_data(loader, 1, samples_num, 1)  # 随机取样本
    inputs = inputs.cuda()
    targets = targets.cuda()

    # 直接分两个组，di dj
    num_group = 2
    N = inputs.shape[0]
    equal_parts = N // num_group
    # 计算 di 训练前 di dj 的损失值
    i = 0
    _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
    before_loss_i = criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts])
    i = 1
    _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
    before_loss_j = criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts])
    # 使用 di 训练
    i = 0
    _optim.zero_grad()
    _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
    _loss = criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts])
    _loss.backward()
    _optim.step()
    # 计算 di 训练后 di dj 的损失值
    i = 0
    _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
    after_loss_i = criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts])
    i = 1
    _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
    after_loss_j = criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts])

    writer.add_scalar('train/delta_loss_%d' % samples_num, before_loss_i-after_loss_i, epoch)
    writer.add_scalar('train/coupling_loss_%d' % samples_num, before_loss_j-after_loss_j, epoch)


def loss_label(net, loader, criterion, scheduler, epoch, writer):

    net = copy.deepcopy(net).train()
    net.zero_grad()
    _optim = optim.SGD(net.parameters(), lr=scheduler.get_last_lr()[0], momentum=0.9, weight_decay=0.0005)

    inputs, targets = fetch_data(loader, 10, 10, mode=0, dm=0)  # 每个样本取10个
    inputs = inputs.cuda()
    targets = targets.cuda()

    # 总样本训练前损失
    _outputs = net.forward(inputs)
    _befloss = criterion(_outputs, targets)
    before_loss = []
    after_loss = []
    # 直接分10个组
    num_group = 10
    N = inputs.shape[0]
    equal_parts = N // num_group
    # 计算训练前每个标签的损失值
    for i in range(num_group):
        _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
        before_loss.append(criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts]).cpu().detach().numpy())
    # 训练
    _optim.zero_grad()
    _outputs = net.forward(inputs)
    _loss = criterion(_outputs, targets)
    _loss.backward()
    _optim.step()
    # 整个损失下降
    _outputs = net.forward(inputs)
    _dloss = _befloss - criterion(_outputs, targets)
    # 计算训练后损失值
    for i in range(num_group):
        _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
        after_loss.append(criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts]).cpu().detach().numpy())

    # 计算每个类别的均值方差
    _delta = [before_loss[i] - after_loss[i] for i in range(num_group)]
    _mean = np.mean(_delta)
    _var = np.var(_delta)

    writer.add_scalar('train/delta_loss', _dloss, epoch)
    writer.add_scalar('train/label_mean_loss', _mean, epoch)
    writer.add_scalar('train/label_var_loss', _var, epoch)


def loss_unit(net, loader, criterion, scheduler, epoch, writer):

    net = copy.deepcopy(net).train()
    net.zero_grad()
    _optim = optim.SGD(net.parameters(), lr=scheduler.get_last_lr()[0], momentum=0.9, weight_decay=0.0005)

    # 直接分组
    num_group = 20

    inputs, targets = fetch_data(loader, 1, num_group, 1)  # 随机取
    inputs = inputs.cuda()
    targets = targets.cuda()

    # 总样本训练前损失
    _outputs = net.forward(inputs)
    _befloss = criterion(_outputs, targets)
    before_loss = []
    after_loss = []
    # 计算训练前每个样本的损失值
    N = inputs.shape[0]
    equal_parts = N // num_group
    for i in range(num_group):
        _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
        before_loss.append(criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts]).cpu().detach().numpy())
    # 训练
    _optim.zero_grad()
    _outputs = net.forward(inputs)
    _loss = criterion(_outputs, targets)
    _loss.backward()
    _optim.step()
    # 整个损失下降
    _outputs = net.forward(inputs)
    _dloss = _befloss - criterion(_outputs, targets)
    # 计算训练后损失值
    for i in range(num_group):
        _outputs = net.forward(inputs[i * equal_parts:(i + 1) * equal_parts])
        after_loss.append(criterion(_outputs, targets[i * equal_parts:(i + 1) * equal_parts]).cpu().detach().numpy())

    # 计算每个样本的均值方差
    _delta = [before_loss[i] - after_loss[i] for i in range(num_group)]
    _mean = np.mean(_delta)
    _var = np.var(_delta)

    writer.add_scalar('train/unit_delta_loss', _dloss, epoch)
    writer.add_scalar('train/unit_mean_loss', _mean, epoch)
    writer.add_scalar('train/unit_var_loss', _var, epoch)


def train_before(net, loader, criterion, epoch, writer, reg=None):
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(loader):
        inputs, targets = inputs.cuda(), targets.cuda()
        outputs = net(inputs)
        loss = criterion(outputs, targets) + reg(net)
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    print('- Loss before training: %.3f | Acc: %.3f%% (%d/%d)' % (train_loss/(batch_idx+1), 100.*correct/total, correct, total))
    if writer:
        writer.add_scalar('train/loss', train_loss / (batch_idx + 1), epoch)
        writer.add_scalar('train/acc', 100. * correct / total, epoch)


class WeightNormWriter:
    def __init__(self):
        super(WeightNormWriter, self).__init__()
        self.w_before = 0
        self.w_after = 0

    def train_before(self, model):
        self.w_before = conv2d_norm(model)

    def train_after(self, model, writer, epoch):
        self.w_after = conv2d_norm(model)
        writer.add_scalar('train/conv2d_norm_change', self.w_before - self.w_after, epoch)


def train_once(mb, trainloader, testloader, config, writer, logger, pretrain=None, lr_mode='cosine', optim_mode='SGD'):
    net = mb.model
    learning_rate = config.learning_rate
    weight_decay = config.weight_decay
    num_epochs = config.epoch
    criterion = nn.CrossEntropyLoss()
    if 'SGD' in optim_mode:
        nesterov = True if 'nest' in optim_mode else False
        optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=weight_decay, nesterov=nesterov)
    elif 'Adam' in optim_mode:
        optimizer = optim.Adam(net.parameters(), lr=learning_rate, weight_decay=weight_decay)
    if lr_mode == 'cosine':
        # lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=0.001)
        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
    elif 'preset' in lr_mode:
        lr_schedule = {0: learning_rate,
                       int(num_epochs * 0.5): learning_rate * 0.1,
                       int(num_epochs * 0.75): learning_rate * 0.01}
        lr_scheduler = PresetLRScheduler(lr_schedule)
    else:
        print('===!!!=== Wrong learning rate decay setting! ===!!!===')
        exit()

    # 用自定义的正则化 - 未解之谜：
    reg = Regularization(weight_decay, 2).cuda()
    optimizer.param_groups[-1]['weight_decay'] = 0

    # 记录每轮的权重范数值变化
    wnw = WeightNormWriter()

    train_before(net, trainloader, criterion, -1, writer, reg)

    print_inf = ''
    best_epoch = 0
    # if pretrain:
    #     best_acc = pretrain['acc']
    #     continue_epoch = pretrain['epoch']
    # else:
    best_acc = 0
    continue_epoch = -1
    for epoch in range(num_epochs):
        if epoch > continue_epoch:  # 其他时间电表空转

            # 临时修改L2参数
            # if epoch == 60:
            #     optimizer.param_groups[-1]['weight_decay'] = 0.0005

            # loss_coupling(net, trainloader, criterion, lr_scheduler, epoch, writer, 256)
            # writer_exp(net, trainloader, criterion, writer, epoch, 128, 'random', True)  # test mode下更好
            # writer_exp(net, trainloader, criterion, writer, epoch, 128, 'random', False)

            # loss_label(net, trainloader, criterion, lr_scheduler, epoch, writer)
            # loss_unit(net, trainloader, criterion, lr_scheduler, epoch, writer)

            wnw.train_before(net)
            train(net, trainloader, optimizer, criterion, epoch, writer, lr_mode, reg)
            test_acc = test(net, testloader, criterion, epoch, writer)
            wnw.train_after(net, writer, epoch)

            if test_acc > best_acc and epoch > 5:
                # print('Saving..')
                state = {
                    'net': net,
                    'acc': test_acc,
                    'epoch': epoch,
                    'args': config,
                    'mask': mb.masks,
                    # 'ratio': mb.get_ratio_at_each_layer()
                }
                path = os.path.join(config.checkpoint_dir, 'train_%s_best.pth.tar' % config.exp_name)
                torch.save(state, path)
                best_acc = test_acc
                best_epoch = epoch
        if lr_mode == 'cosine':
            _lr = lr_scheduler.get_last_lr()
            lr_scheduler.step()

            # 临时修改学习率
            # if epoch > 9:
            #     for param_group in optimizer.param_groups:
            #         param_group['lr'] = 0.9 * param_group['lr']

        else:
            _lr = lr_scheduler.get_lr(optimizer)
            lr_scheduler(optimizer, epoch)
        if writer:
            writer.add_scalar('train/lr', _lr, epoch)

    logger.info('best acc: %.4f, epoch: %d' % (best_acc, best_epoch))
    return 'best acc: %.4f, epoch: %d\n' % (best_acc, best_epoch), print_inf

